# -*- coding: utf-8 -*-            
# @Time    : 2025/5/7 11:31
# @FileName: test.py
# @Software: PyCharm
# 导入必要的库
import torch  # PyTorch深度学习框架
import numpy as np  # 数值计算库
from sklearn.datasets import load_boston  # 波士顿房价数据集
from sklearn.preprocessing import MinMaxScaler  # 数据归一化
from sklearn.model_selection import train_test_split  # 数据集划分

# 加载波士顿房价数据集
boston = load_boston()
# 获取特征数据（13个特征）
data = boston.data
# 获取目标值（房价）并重塑为列向量
target = boston.target.reshape(-1, 1)

# 创建数据归一化器
scaler_data = MinMaxScaler()  # 用于特征数据
scaler_target = MinMaxScaler()  # 用于目标值
# 对数据进行归一化（将所有特征缩放到[0,1]范围）
data = scaler_data.fit_transform(data)
# 对目标值进行归一化
target = scaler_target.fit_transform(target)

# 设置时间窗口大小（用过去7天的数据预测下一天）
c = 7
# 准备输入序列(x)和输出序列(y)
x = []
y = []
# 创建时间序列数据
for i in range(len(data) - c):
    # 获取从i开始的c个时间步的特征数据作为输入
    x_data = data[i:i + c]
    # 获取第i+c个时间步的目标值作为输出
    y_data = target[i + c]
    x.append(x_data)
    y.append(y_data)

# 将numpy数组转换为PyTorch张量
x = torch.Tensor(x)
y = torch.Tensor(y)

# 划分训练集和测试集（不进行shuffle以保持时间序列顺序）
train_x, test_x, train_y, test_y = train_test_split(x, y, shuffle=False)
# 打印训练集的形状（样本数, 时间步长, 特征数）
print(train_x.shape)

# 定义GRU模型类
class GRUModel(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        # 创建GRU层
        # input_size: 输入特征维度（13个特征）
        # hidden_size: 隐藏层维度（50个单元）
        # batch_first: 输入张量的第一个维度是batch
        self.gru = torch.nn.GRU(input_size=x.shape[2], hidden_size=50, batch_first=True)
        # 创建全连接层，将GRU输出映射到1维（预测房价）
        self.fc = torch.nn.Linear(in_features=50, out_features=1)

    def forward(self, x):
        # 前向传播
        # gru返回: output (所有时间步的隐藏状态), h_n (最后一个时间步的隐藏状态)
        x, _ = self.gru(x)
        # 取最后一个时间步的输出，通过全连接层得到预测值
        out = self.fc(x[:, -1, :])
        return out

# 实例化模型
model = GRUModel()
# 定义损失函数（均方误差）
loss_fn = torch.nn.MSELoss()
# 定义优化器（Adam优化算法，学习率0.01）
op = torch.optim.Adam(model.parameters(), lr=0.01)
# 存储损失值的列表
loss_list = []

# 训练循环（1000个epoch）
for i in range(1000):
    # 清空梯度
    op.zero_grad()
    # 前向传播，获取预测值
    h = model(train_x)
    # 计算损失
    loss = loss_fn(h, train_y)
    # 反向传播
    loss.backward()
    # 注释掉的代码：记录损失值
    # loss_list.append(loss.item())
    # 更新参数
    op.step()

# 在测试集上进行预测
pre = model(test_x).reshape(-1)  # 将预测结果重塑为一维数组
# 打印预测结果
print(pre)

# 导入matplotlib用于可视化
import matplotlib.pyplot as plt

# 绘制预测值和真实值的对比图
plt.plot(pre.data.numpy(), c='r', label='Prediction')  # 红色线表示预测值
plt.plot(test_y.data.numpy(), c='g', label='Actual')   # 绿色线表示真实值
plt.legend()  # 显示图例
plt.show()  # 显示图形